import cv2
import numpy as np
import torch
from PIL import Image
import torchvision.transforms as T
from pathlib import Path

class SmartSegmentation:
    """智能分割工具类，提供多种分割算法支持"""
    
    def __init__(self, device=None):
        """
        初始化智能分割工具
        
        Args:
            device: 使用的设备，如果未指定，则自动选择
        """
        # 设置设备
        if device is None:
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        else:
            self.device = device
        
        # 加载模型状态
        self.model_loaded = False
        self.model = None
        self.transform = None
    
    def load_model(self, model_type="deeplabv3"):
        """
        加载分割模型
        
        Args:
            model_type (str): 模型类型，支持 "deeplabv3", "fcn"
            
        Returns:
            bool: 是否成功加载模型
        """
        try:
            if model_type == "deeplabv3":
                # 加载DeepLabV3模型
                self.model = torch.hub.load('pytorch/vision:v0.10.0', 
                                         'deeplabv3_resnet50', 
                                         pretrained=True)
            elif model_type == "fcn":
                # 加载FCN模型
                self.model = torch.hub.load('pytorch/vision:v0.10.0', 
                                         'fcn_resnet50', 
                                         pretrained=True)
            else:
                print(f"不支持的模型类型: {model_type}")
                return False
            
            # 设置模型评估模式
            self.model.eval()
            self.model.to(self.device)
            
            # 设置转换
            self.transform = T.Compose([
                T.ToTensor(),
                T.Normalize(mean=[0.485, 0.456, 0.406], 
                            std=[0.229, 0.224, 0.225]),
            ])
            
            self.model_loaded = True
            return True
        
        except Exception as e:
            print(f"加载模型时出错: {e}")
            self.model_loaded = False
            return False
    
    def segment_image(self, image, threshold=0.7):
        """
        对图像进行语义分割
        
        Args:
            image: OpenCV格式的图像(BGR)或图像路径
            threshold: 分割阈值
            
        Returns:
            mask: 分割掩码，每个像素为分类ID
            polygons: 找到的多边形轮廓列表
        """
        # 检查模型是否已加载
        if not self.model_loaded:
            loaded = self.load_model()
            if not loaded:
                return None, []
        
        # 加载并处理图像
        if isinstance(image, str) or isinstance(image, Path):
            # 从文件加载图像
            img = cv2.imread(str(image))
            if img is None:
                return None, []
        else:
            img = image
        
        # 转换为RGB(PIL格式)
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        input_image = Image.fromarray(img_rgb)
        
        # 使用预训练模型进行预测
        input_tensor = self.transform(input_image)
        input_batch = input_tensor.unsqueeze(0).to(self.device)
        
        with torch.no_grad():
            output = self.model(input_batch)['out'][0]
        
        # 处理输出
        output_predictions = output.argmax(0).byte().cpu().numpy()
        
        # 提取非背景类别的掩码
        mask = np.zeros_like(output_predictions)
        mask[output_predictions > 0] = 255
        
        # 查找轮廓
        contours, _ = cv2.findContours(
            mask.astype(np.uint8), 
            cv2.RETR_EXTERNAL, 
            cv2.CHAIN_APPROX_SIMPLE
        )
        
        # 转换轮廓为多边形
        polygons = []
        for contour in contours:
            # 简化轮廓为多边形
            epsilon = 0.005 * cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, epsilon, True)
            
            # 如果多边形点数足够且面积不太小，则加入结果
            if len(approx) >= 3 and cv2.contourArea(approx) > 100:
                # 提取多边形点
                points = []
                for point in approx:
                    x, y = point[0]
                    points.append({"x": float(x), "y": float(y)})
                
                polygons.append(points)
        
        return mask, polygons
    
    def interactive_segment(self, image, points, labels):
        """
        基于用户交互点进行分割
        
        Args:
            image: OpenCV格式的图像(BGR)或图像路径
            points: 交互点列表，每点为(x, y)元组
            labels: 点的标签列表，1表示前景，0表示背景
            
        Returns:
            mask: 分割掩码
            polygon: 找到的多边形轮廓
        """
        # 加载并处理图像
        if isinstance(image, str) or isinstance(image, Path):
            # 从文件加载图像
            img = cv2.imread(str(image))
            if img is None:
                return None, []
        else:
            img = image
        
        # 创建蒙版和输入点
        mask = np.zeros(img.shape[:2], dtype=np.uint8)
        
        # 转换点格式
        fg_points = []
        bg_points = []
        
        for i, (point, label) in enumerate(zip(points, labels)):
            if label == 1:
                fg_points.append(point)
            else:
                bg_points.append(point)
                
        # 如果没有足够的前景/背景点，则返回
        if not fg_points:
            return None, []
        
        # 创建grabCut需要的数据
        gc_mask = np.zeros(img.shape[:2], dtype=np.uint8)
        bgd_model = np.zeros((1, 65), dtype=np.float64)
        fgd_model = np.zeros((1, 65), dtype=np.float64)
        
        # 使用foreground points设置确定的前景
        for point in fg_points:
            x, y = int(point[0]), int(point[1])
            cv2.circle(gc_mask, (x, y), 3, cv2.GC_FGD, -1)
        
        # 使用background points设置确定的背景
        for point in bg_points:
            x, y = int(point[0]), int(point[1])
            cv2.circle(gc_mask, (x, y), 3, cv2.GC_BGD, -1)
        
        # 设置矩形外的区域为确定背景
        rect = (0, 0, img.shape[1], img.shape[0])
        
        # 运行GrabCut算法
        cv2.grabCut(
            img, gc_mask, rect, bgd_model, fgd_model, 
            5, cv2.GC_INIT_WITH_MASK
        )
        
        # 生成二值掩码
        mask_binary = np.zeros_like(gc_mask)
        mask_binary[gc_mask == cv2.GC_FGD] = 255
        mask_binary[gc_mask == cv2.GC_PR_FGD] = 255
        
        # 查找轮廓
        contours, _ = cv2.findContours(
            mask_binary, 
            cv2.RETR_EXTERNAL, 
            cv2.CHAIN_APPROX_SIMPLE
        )
        
        # 转换轮廓为多边形
        polygons = []
        for contour in contours:
            # 简化轮廓为多边形
            epsilon = 0.005 * cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, epsilon, True)
            
            # 如果多边形点数足够且面积不太小，则加入结果
            if len(approx) >= 3 and cv2.contourArea(approx) > 100:
                # 提取多边形点
                points = []
                for point in approx:
                    x, y = point[0]
                    points.append({"x": float(x), "y": float(y)})
                
                polygons.append(points)
        
        return mask_binary, polygons 